import datetime
import functools
import logging
import pandas as pd

from absl import app
from tornado.ioloop import IOLoop, PeriodicCallback

import coin.flow.subscriber as flow_subscriber
from coin.flow.simple_queue import SimpleQueue
from coin.base.datetime_util import to_datetime
from coin.base.timestamp import get_timestamp

from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.feed.fastfeed.native_wrapper import run_from_feed_cache

import experimental.prophet as prophet
from coin.exchange.okex.kr_rest.product import OkexProduct


class Test:
  def __init__(self):
    self._vars = []
    self._var_names = []
    self._live_begin_ts = None

    self._lookback = 10
    self._rows = []

  def model(self, live):
    ts = prophet.global_timestamp()
    if not live:
      # Timer does not work in live.
      timer_1s = prophet.timer('1s')

    product = OkexProduct.FromStr('BTC-USDT')
    feed = prophet.fastfeed_coin(product, 'Spot.Okex.v3')
    trade = prophet.trade(feed)

    tq = prophet.time_moving_window(trade.qty, '1s').sum().named_ref('trade_qty')
    tv = prophet.time_moving_window(trade.qty * trade.price, '1s').sum().named_ref('trade_volume')
    vwap = (tv / tq).named_ref('vwap')

    vars = [ts, tq, tv, vwap]
    return vars

  def matrix(self):
    ts = prophet.global_timestamp()
    timer_1s = prophet.timer('1s')
    vars = self.model(False)
    with prophet.control_if(ts % 10**9 == 0):
      aggregator = prophet.scalar_value_aggregator(vars)
    return aggregator

  def on_eval(self, graph, live=False):
    ts = graph.get_named_variable('global_timestamp').value
    if live and ts < self._live_begin_ts:
      return

    if not (live or ts % 10**9 == 0):
      return

    row = [var.value for var in self._vars]
    self._rows.append(row)

    df = pd.DataFrame(self._rows[-self._lookback:], columns=self._var_names)
    print(to_datetime(ts))
    print(df)

  def dump_matrix(self):
    df = prophet.run_from_fastfeed_multiprocess(self.matrix,
                                                datetime.date(2019, 8, 1),
                                                datetime.datetime(2019, 8, 1, 2),
                                                machine='feed-01.ap-northeast-1.aws',
                                                use_run_cache=False,
                                                duration_after='1m')
    print(df)

  def run_from_fastfeed(self):
    graph = prophet.Graph()
    with prophet.as_default_graph(graph):
      self._vars = self.model(False)
      self._var_names = [var.name for var in self._vars]

    machine = 'feed-01.ap-northeast-1.aws'
    adapter = prophet.FastfeedGraphAdapter(graph, machine, eval_callback=self.on_eval)
    feed_sub_req = adapter.get_feed_sub_requests()

    run_from_feed_cache(feed_sub_req,
                        adapter.on_book_reset,
                        datetime.date(2019, 8, 1),
                        datetime.datetime(2019, 8, 1, 2),
                        machine=machine)

  def run_live(self):
    ioloop = IOLoop.current()
    queue = SimpleQueue('queue')
    flow_sub = flow_subscriber.from_queue(queue)

    graph = prophet.Graph()
    with prophet.as_default_graph(graph):
      self._vars = self.model(True)
      self._var_names = [var.name for var in self._vars]

    adapter = prophet.FastfeedGraphAdapter(graph, '')
    feed_sub_req = adapter.get_feed_sub_requests()

    feed_subsystem = MultiFeedSubsystem()
    feed_subsystem.subscribe_fastfeed(feed_sub_req, adapter.on_book_reset)
    feed_subsystem.start(ioloop, queue, flow_sub)

    interval_trigger = PeriodicCallback(functools.partial(self.on_eval, graph, True), 1000)
    interval_trigger.start()
    self._live_begin_ts = get_timestamp()

    try:
      ioloop.start()
      return 0
    except KeyboardInterrupt:
      return 1


def main(argv):
  tester = Test()

  if len(argv) == 1:
    tester.run_live()
  elif argv[1] == 'live':
    tester.run_live()
  elif argv[1] == 'archive':
    tester.run_from_fastfeed()
  elif argv[1] == 'dump':
    tester.dump_matrix()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG')
  app.run(main)
